import warnings
import numpy as np
import os
from tqdm import tqdm
from PIL import Image
from io import BytesIO
from lang_sam import LangSAM
import time
import requests

# 图像库路径
image_path = "/home/t/workspace/XMem/workspace/177661140-f690156b-1775-4cd7-acd7-1738a5c92f30/images"
text_prompt = "animal"

def download_image(url):
    response = requests.get(url)
    response.raise_for_status()
    return Image.open(BytesIO(response.content)).convert("RGB")

def save_mask(mask_np, filename):
    mask_image = Image.fromarray((mask_np * 255).astype(np.uint8))
    mask_image.save(filename)

def main():
    # 抑制警告消息
    warnings.filterwarnings("ignore")

    model = LangSAM()

    # 获取所有图像文件
    image_files = [f for f in os.listdir(image_path) if f.endswith('.jpg')]
    num_images = len(image_files)

    # 使用 tqdm 进度条
    pbar = tqdm(total=num_images, unit='image')

    for image_file in image_files:
        try:
            image_pil = Image.open(os.path.join(image_path, image_file)).convert("RGB")

            # 测量时间消耗
            start = time.time()
            masks, boxes, phrases, logits = model.predict(image_pil, text_prompt, box_threshold=0.5)
            end = time.time()
            print(f"Time consumption for {image_file}: {end - start:.2f}s")

            # 创建文件夹存储mask
            mask_dir = os.path.splitext(image_file)[0]
            os.makedirs(mask_dir, exist_ok=True)

            for i, mask in enumerate(masks):
                mask_np = mask.squeeze().cpu().numpy()
                mask_path = os.path.join(mask_dir, f"mask_{i+1}.png")
                save_mask(mask_np, mask_path)

            pbar.update(1)

        except (requests.exceptions.RequestException, IOError) as e:
            print(f"Error: {e}")
            pbar.update(1)

    pbar.close()

if __name__ == "__main__":
    main()
